import os
import sys
import argparse
import pandas as pd

def convert_test_to_complex_format(input_file, output_file):
    """
    将test.csv格式转换为complex_cell_data.csv格式
    
    转换规则：
    - 将复合列名（如xxxx_0.001）拆分为compound和concentration_um
    - 创建well_id列，格式为：cell_line_compound_concentration_um_Rep1
    - viability_percentage值从原始数据中获取
    """
    print(f"正在读取输入文件: {input_file}")
    
    # 读取输入文件
    df = pd.read_csv(input_file)
    
    print(f"原始数据形状: {df.shape}")
    print(f"原始数据列名: {list(df.columns)}")
    
    # 准备输出数据列表
    output_data = []
    
    # 处理每一行数据
    for row in df.iterrows():
        cell_line = row['cell_line']
        
        # 遍历所有复合列名（排除cell_line列）
        for col_name in df.columns[1:]:  # 跳过第一列cell_line
            # 分割复合列名：compound_concentration_um
            if '_' in col_name:
                # 尝试分割为compound和concentration
                parts = col_name.split('_')
                if len(parts) >= 2:
                    # 处理可能包含下划线的化合物名称
                    compound = '_'.join(parts[:-1])  # 化合物名称可能包含下划线
                    concentration_um = parts[-1]    # 最后一部分是浓度
                    
                    # 确保浓度是数字格式
                    try:
                        concentration_um = float(concentration_um)
                        # 获取viability_percentage值并转换为百分比形式
                        viability_value = row[col_name]
                        # 将小数值转换为百分比（乘以100），保留原始值的精度
                        viability_percentage = viability_value * 100 if viability_value < 2 else viability_value
                        
                        # 创建well_id
                        well_id = f"{cell_line}_{compound}_{concentration_um}_Rep1"
                        
                        # 添加到输出数据
                        output_data.append({
                            'cell_line': cell_line,
                            'compound': compound,
                            'concentration_um': concentration_um,
                            'well_id': well_id,
                            'viability_percentage': viability_percentage
                        })
                    except ValueError:
                        print(f"警告：无法将浓度值 '{parts[-1]}' 转换为数字，跳过列 '{col_name}'")
            else:
                print(f"警告：列名 '{col_name}' 不包含下划线，无法拆分，跳过")
    
    # 创建输出DataFrame
    output_df = pd.DataFrame(output_data)
    
    print(f"转换后的数据形状: {output_df.shape}")
    print(f"转换后的列名: {list(output_df.columns)}")
    
    # 保存到输出文件
    output_df.to_csv(output_file, index=False)
    print(f"数据已保存至: {output_file}")
    
    # 显示前几行数据以验证
    print("\n转换后的数据预览:")
    print(output_df.head())
    
    # 显示唯一的cell_line和compound
    print(f"\n唯一的cell_line数量: {output_df['cell_line'].nunique()}")
    print(f"唯一的compound数量: {output_df['compound'].nunique()}")
    print(f"浓度范围: {output_df['concentration_um'].min()} - {output_df['concentration_um'].max()} uM")

if __name__ == "__main__":
    # 创建参数解析器
    parser = argparse.ArgumentParser(description='将test.csv格式转换为complex_cell_data.csv格式')
    
    # 添加命令行参数
    parser.add_argument(
        '--input', 
        type=str, 
        default='test.csv',
        help='输入文件路径（默认为test.csv）'
    )
    parser.add_argument(
        '--output', 
        type=str, 
        default='complex_cell_data.csv',
        help='输出文件路径（默认为complex_cell_data.csv）'
    )
    
    # 解析参数
    args = parser.parse_args()
    
    # 获取输入和输出文件路径
    input_file = args.input
    output_file = args.output
    
    # 检查输入文件是否存在
    if not os.path.exists(input_file):
        print(f"错误：输入文件 '{input_file}' 不存在")
        sys.exit(1)
    
    # 创建输出目录（如果不存在）
    output_dir = os.path.dirname(output_file)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 执行转换
    convert_test_to_complex_format(input_file, output_file)
    print("\n数据转换完成！")